import { NextRequest, NextResponse } from "next/server";

import { ChatOpenAI } from "@langchain/openai";
import { PromptTemplate } from "@langchain/core/prompts";
import { StructuredOutputParser } from "langchain/output_parsers";
import { RunnableSequence } from "@langchain/core/runnables";
import { z } from "zod";

export const runtime = "edge";

const TEMPLATE = `Extract the requested fields from the input.

The field "entity" refers to the first mentioned entity in the input.

Input:

{input}

{format_instructions}`;


export async function POST(req: NextRequest) {
  try {
    const body = await req.json();
    const messages = body.messages ?? [];
    const currentMessageContent = messages[messages.length - 1].content;

    const prompt = PromptTemplate.fromTemplate(TEMPLATE);

    const model = new ChatOpenAI({
      temperature: 0.8,
      modelName: "gpt-3.5-turbo",
    });

    const schema = z.object({
      tone: z
        .enum(["positive", "negative", "neutral"])
        .describe("The overall tone of the input"),
      entity: z.string().describe("The entity mentioned in the input"),
      word_count: z.number().describe("The number of words in the input"),
      chat_response: z.string().describe("A response to the human's input"),
      final_punctuation: z
        .optional(z.string())
        .describe("The final punctuation mark in the input, if any."),
    });

    const parser = StructuredOutputParser.fromZodSchema(schema);
    const chain = RunnableSequence.from([
      prompt,
      model,
      parser
    ]);

    const result = await chain.invoke({
      input: currentMessageContent,
      format_instructions: parser.getFormatInstructions(),
    });

    return NextResponse.json(result, { status: 200 });
  } catch (e: any) {
    return NextResponse.json({ error: e.message }, { status: e.status ?? 500 });
  }
}

// import { NextRequest, NextResponse } from "next/server";

// import { z } from "zod";
// import { zodToJsonSchema } from "zod-to-json-schema";

// import { ChatOpenAI } from "@langchain/openai";
// import { PromptTemplate } from "@langchain/core/prompts";
// import { JsonOutputFunctionsParser } from "langchain/output_parsers";

// export const runtime = "edge";

// const TEMPLATE = `Extract the requested fields from the input.

// The field "entity" refers to the first mentioned entity in the input.

// Input:

// {input}`;

// export async function POST(req: NextRequest) {
//   try {
//     const body = await req.json();
//     const messages = body.messages ?? [];
//     const currentMessageContent = messages[messages.length - 1].content;

//     const prompt = PromptTemplate.fromTemplate(TEMPLATE);

//     const model = new ChatOpenAI({
//       temperature: 0.8,
//       modelName: "gpt-4",
//     });

//     const schema = z.object({
//       tone: z
//         .enum(["positive", "negative", "neutral"])
//         .describe("The overall tone of the input"),
//       entity: z.string().describe("The entity mentioned in the input"),
//       word_count: z.number().describe("The number of words in the input"),
//       chat_response: z.string().describe("A response to the human's input"),
//       final_punctuation: z
//         .optional(z.string())
//         .describe("The final punctuation mark in the input, if any."),
//     });

//     const functionCallingModel = model.bind({
//       functions: [
//         {
//           name: "output_formatter",
//           description: "Should always be used to properly format output",
//           parameters: zodToJsonSchema(schema),
//         },
//       ],
//       function_call: { name: "output_formatter" },
//     });
//     const chain = prompt
//       .pipe(functionCallingModel)
//       .pipe(new JsonOutputFunctionsParser());

//     const result = await chain.invoke({
//       input: currentMessageContent,
//     });

//     return NextResponse.json(result, { status: 200 });
//   } catch (e: any) {
//     return NextResponse.json({ error: e.message }, { status: e.status ?? 500 });
//   }
// }